import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Random;

public class KMeans {
    private int k; // 簇的个数
    private int maxIterations; // 最大迭代次数
    private double[][] data; // 数据集
    private double[][] centroids; // 簇中心点
    private int[] labels; // 每个数据点所属的簇
    private int numExamples; // 数据集大小
    private int numFeatures; // 特征数量
    private int clusterCapacity; // 簇上点的个数不能超过宿舍容量

    public KMeans(int k, int maxIterations, double[][] data, int clusterCapacity) {
        this.k = k;
        this.maxIterations = maxIterations;
        this.data = data;
        this.numExamples = data.length;
        this.numFeatures = data[0].length;
        this.centroids = new double[k][numFeatures];
        this.labels = new int[numExamples];
        Arrays.fill(labels, -1);
        this.clusterCapacity = clusterCapacity;
    }

    public void run() {
        initCentroids();
        int iterations = 0;
        boolean shouldStop = false;

        while (iterations < maxIterations && !shouldStop) {
            shouldStop = true;

            // Assign points to closest centroids
            for (int i = 0; i < numExamples; i++) {
                int closestCluster = getClosestCluster(data[i]);
                if (labels[i] != closestCluster) {
                    labels[i] = closestCluster;
                    shouldStop = false;
                }
            }

            // Update centroids
            for (int i = 0; i < k; i++) {
                int numPoints = 0;
                double[] newCentroid = new double[numFeatures];
                for (int j = 0; j < numExamples; j++) {
                    if (labels[j] == i) {
                        for (int f = 0; f < numFeatures; f++) {
                            newCentroid[f] += data[j][f];
                        }
                        numPoints++;
                    }
                }

                if (numPoints > 0) {
                    for (int f = 0; f < numFeatures; f++) {
                        newCentroid[f] /= numPoints;
                    }
                    centroids[i] = newCentroid;
                }
            }

            iterations++;
        }

        assignClustersToDormitories();
    }

    private void initCentroids() {
        Random rand = new Random();
        for (int i = 0; i < k; i++) {
            boolean newCentroid = true;
            while (newCentroid) {
                int idx = rand.nextInt(numExamples);
                centroids[i] = data[idx];
                newCentroid = false;
                for (int j = 0; j < i; j++) {
                    if (Arrays.equals(centroids[i], centroids[j])) {
                        newCentroid = true;
                        break;
                    }
                }
            }
        }
    }

    private int getClosestCluster(double[] example) {
        double minDist = Double.MAX_VALUE;
        int closestCluster = -1;
        for (int i = 0; i < k; i++) {
            double dist = euclideanDistance(example, centroids[i]);
            if (dist < minDist) {
                minDist = dist;
                closestCluster = i;
            }
        }
        return closestCluster;
    }

    private double euclideanDistance(double[] a, double[] b) {
        double sum = 0.0;
        for (int i = 0; i < numFeatures; i++) {
            sum += Math.pow(a[i] - b[i], 2);
        }
        return Math.sqrt(sum);
    }

    private void assignClustersToDormitories() {
        int totalNumRooms = (int) Math.ceil((double) numExamples / clusterCapacity);
        int[] numRooms = new int[k];
        int[] numRoomsPerCluster = new int[k];
        int[] clusterSizes = new int[k];

        for (int i = 0; i < numExamples; i++) {
            int cluster = labels[i];
            clusterSizes[cluster]++;
            numRoomsPerCluster[cluster] = (int) Math.ceil((double) clusterSizes[cluster] / clusterCapacity);
            numRooms[cluster] = Math.min(numRoomsPerCluster[cluster], totalNumRooms); // 簇上点的个数不能超过宿舍容量
        }

        List<Integer>[] clusters = new List[k];
        for (int i = 0; i < k; i++) {
            clusters[i] = new ArrayList<>();
        }

        for (int i = 0; i < numExamples; i++) {
            int cluster = labels[i];
            int room = clusters[cluster].size() % numRooms[cluster] + 1;
            clusters[cluster].add(room);
        }

        for (int i = 0; i < k; i++) {
            System.out.println("Dormitory " + (i + 1) + ":");
            for (int j = 0; j < clusters[i].size(); j++) {
                System.out.print("  Room " + (j + 1) + ": ");
                for (int k = 0; k < numExamples; k++) {
                    if (labels[k] == i && clusters[i].get(j).equals(clusters[i].get((k+1)%clusters[i].size()))) {
                        System.out.print((k + 1));
                        if (k < clusters[i].size() - 1) {
                            System.out.print(", ");
                        }
                    }
                }
                System.out.println();
            }
        }
    }
}